import sys
import os
current_file_path = os.path.abspath(__file__)
project_root = os.path.dirname(os.path.dirname(current_file_path))
sys.path.append(project_root)
from typing import List, TypedDict
import time
from pydantic import BaseModel, Field
from collections import OrderedDict
import asyncio
from typing import TypedDict, List, Literal, Any
from llm import create_llm
from report.chromedb.index import ChromaDBManager
from langchain_core.prompts import ChatPromptTemplate
from pathlib import Path
from langgraph.prebuilt import create_react_agent
from tools.query_tools import query_tools
from tools.count_tools import count_tools
import logging
import httpx
httpx.get("https://api.deepseek.com", timeout=5)  # 测试基础连通性
logger = logging.getLogger(__name__)  # 在需要的地方创建logger实例
# 在文件开头添加日志配置（可选但推荐）
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)





db_manager = ChromaDBManager()


# 匹配模板
def match_template(input: str) -> bool:
    results = db_manager.search_data(
        query=input,
        filter_fields={"category": "报告模板"},
        k=1,
        similarity_threshold=0.5
    )
    # 匹配模板
    return results
# 匹配提示词
def match_prompt(input: str) -> bool:
    results = db_manager.search_data(
        query=input,
        filter_fields={"category": "提示词"},
        k=1,
        similarity_threshold=0.6
    )
    # 匹配模板
    return results




# 可获取数据
getData = """
1.学生信息（[单个/某班级（三年二班）/某年级（3）]学生信息
示例：
    张三学生信息
    三年二班学生信息
    3年级学生信息
2.学生科目成绩（[单个/某班级（三年二班）/某年级（3）]学生单科目或所有科目成绩
示例：
    张三语文成绩
    三年二班语文成绩
    3年级语文成绩
    张三所有科目成绩
    三年二班所有科目成绩
    3年级所有科目成绩
3.学生成绩总分（单个学生总分）
示例：
    张三总分
4.考试分数等级（如：A+: 95-100，A:90-94...）
"""

# 定义AgentState
class AgentState(TypedDict):
    input: str
    output: str
    data: List[dict]
    currentData: List[Any]
    type: str
    templateList: List[str]
    template: str
    demandList: List[str]
    total_retries: int
    error_tip: str
    paragraphList: List[str]



# 生成模板智能体
class GenerateDescription(BaseModel):
    template: str = Field(
        description="根据用户要求生成的模板"
    )
generate_prompt =  ChatPromptTemplate.from_messages(
    [
        ("system", f"""
            你是一个模板生成助手，你将根据用户的要求返回一个模板，模板展示的数据只能是以下数据或者可以根据以下数据计算所得的数据。
            当前可以拿到的数据有：
            {getData}
            你要按照以下步骤执行：
            1.分析用户需求是否是一个报告
            2-A.如果是
                1.判断是否为个人报告
                2-a.如果是个人报告
                    1.将第一行的姓名替换报告中的姓名
                    2.将替换后的报告返回
                2-b.如果不是个人报告
                    1.将报告直接返回
            2-B.如果不是
                1.根据用户的需求和可获取数据创建一个模板
            3.将转换或者创建的模板返回
            注：模板后续会进行数据获取填充，不要让返回的模板里有无法获取的数据
        """),
        ("human", "用户需求：{input}")
    ]
)
generate = generate_prompt | create_llm("deepseek", 1, 60).with_structured_output(GenerateDescription)

# 检查美化模板智能体
class InspectDescription(BaseModel):
    template: str = Field(
        description="修改美化后的模板"
    )
inspect_prompt =  ChatPromptTemplate.from_messages(
    [
        ("system", f"""
            你是一个模板优化助手，你将根据用户的需求和现有的模板案例来生成一份新的模板，这份模板的数据必须是可以获取的数据或根据获取的数据计算得出的数据，可以获取的数据有：
            {getData}
            你将按照以下步骤执行：
            1.分析模板案例，将里面的所有示例数据替换为模板描述，如：班级有10名学生，替换为班级有{{{{{{{{班级学生数量}}}}}}}}名学生
            2.分析用户的需求，查看当前模板案例是否符合用户需求，如果不符合，修改模板案例
            3.将修改后的模板使用markdown语法进行美化
            4.将修改后的模板案例返回
            注：模板后续会进行数据获取填充，不要让返回的模板里有无法获取的数据
        """),
        ("human", "模板案例：{template}")
    ]
)
inspect = inspect_prompt | create_llm("deepseek", 1, 60).with_structured_output(InspectDescription)


# 需求分析智能体输出格式化
class VerifyDescription(BaseModel):
    demand: List[str] = Field(
        description="当前需要的数据列表，给数据获取器的信息，必须数据获取器可以获取到的数据"
    )
# 需求分析智能体
verify_prompt =  ChatPromptTemplate.from_messages(
    [
        ("system", f"""
            你是一个模板分析助手，你将根据模板进行分析，返回需要获取的数据列表，交由数据获取器。
            数据获取器可以直接获取到的数据有：
            {getData}
            你将按照以下步骤执行：
            1.判断模板需要哪些数据，数据根据示例进行返回，将示例中关键字（如张三、三年二班等）替换为用户所需后返回，不要直接返回示例，
            2.返回替换后的模板所需数据列表
            注：返回所需数据列表内容必须为数据获取器可以获取到的数据，尽量减少数据获取次数，必须一次返回所有需要的数据，若返回超出数据获取器可以获取的数据，则会触发错误
        """),
        ("human", "模板：{template}")
    ]
)
verify = verify_prompt | create_llm("deepseek", 1, 60).with_structured_output(VerifyDescription)

# 数据获取智能体
query_data_prompt = """
    你是一个数据获取助手，你将按照以下步骤进行数据获取：
    1.根据问题选择对应的工具进行数据获取
    2.获取到数据后，返回给用户
    3.若获取到的数据内容为空，则返回“当前查询暂无数据，请确定问题正确性！”
    默认查询用户询问的所有数据。
"""
query_agent = create_react_agent(
    model=create_llm("glm", 0.3, 60),
    tools=query_tools,
    prompt=query_data_prompt,
)
# 数据压缩智能体
class CompressedData(BaseModel):
    compressed: dict = Field(
        description="压缩后的数据字典，使用最短键名和简化结构",
        examples=[{"s": "张三", "c":80, "m":85}]
    )
    mapping: dict = Field(
        description="键名映射关系（如{'s':'student_name'}），用于后续解压",
        examples=[{"s": "student_name", "c": "chinese_score"}]
    )
    isNone: bool = Field(
        description="数据内容为空，返回True，否则返回False"
    )
compression_prompt =  ChatPromptTemplate.from_messages(
    [
        ("system", """
            你是一个JSON压缩专家，请按以下策略压缩数据：
            1. 【键名缩写】用单个字母替换长字段名（如studentName→s）
            2. 【结构扁平化】将嵌套结构展开为键路径（如address.city→addr_city）
            3. 【数值优化】用Number类型替代字符串数字
            4. 【数据合并】将相似结构数组合并为矩阵（如[80,85,90]→"80,85,90"）
            5. 【映射分离】单独维护键名映射字典
            6. 【token统计】对比压缩前后token量
            7. 【字段顺序标准化】统一按[学号,姓名,成绩]顺序合并数据
            8. 【统一合并策略】同类数据优先使用垂直拆分（数组存储）
            9. 【异常值处理】自动修正字段中的空格/格式问题
            输出格式要求：
            - 优先使用列表等紧凑结构
            - 去除所有JSON格式化空格
            - 保留至少3位小数精度
        """),
        ("human", "待压缩数据：{currentData}")
    ]
)
refine =  compression_prompt | create_llm("glm", 0.1, 30, True).with_structured_output(CompressedData)
# 数据加工输出智能体
process_prompt = """
    你是一个教务系统AI助手，你的工作是进行数据加工输出：
    你将按照以下步骤执行：
    1.分析模板。
    2.根据模板观察缺少那些变量。
    3.使用工具处理现有数据，获取模板所需数据。
    4.使用计算出的数据填充模板。
    5.使用markdown语法进行美化。
    6.将填充后的模板返回给用户（只需要返回填充后的模板，其他多余内容（步骤过程等内容）不要返回，模板无需使用代码块包裹）。
    你是一个数学白痴，所有和计算有关的问题你都不会算，但是你擅长使用计算器，你会用计算器来完成运算，四则运算全部使用计算器完成，禁止直接计算。
    已有数据为压缩数据，mapping为映射关系
"""
process = create_react_agent(
    model = create_llm("deepseek", 0.3, 120),
    tools=count_tools,
    prompt=process_prompt,
)




def main(input: str):

    # 未匹配模板生成节点
    def generate_step(state: AgentState):
        value = generate.invoke({"input": state["input"]})
        print("---------未匹配模板生成执行结果：-", value)
        return {
            "template": value.template,
            "data": [],
            "total_retries": 0,
            "error_tip": ""
        }
    
    # 优化模板节点
    def inspect_step(state: AgentState):
        params = {
            "template": state["template"],
        }
        value = inspect.invoke(params)
        print("---------优化模板执行结果：-", value.template)
        return {
            "template": value.template,
        }

    # 需求分析节点
    def verify_step(state: AgentState):
        params = {
            "template": state["template"],
        }
        print("---------需求分析执行参数：-", params)
        verifyValue = verify.invoke(params)
        print("---------需求分析执行结果：-", verifyValue)
        return {
            "demandList": verifyValue.demand,
        }
    # 并行获取数据节点
    def parallel_query_step(state: AgentState):
        from concurrent.futures import ThreadPoolExecutor, as_completed
        from tenacity import retry, stop_after_attempt, wait_exponential
        import asyncio
        
        success_count = 0
        error_count = 0
        @retry(stop=stop_after_attempt(3), 
            wait=wait_exponential(multiplier=1, min=4, max=10))
        def safe_query(demand):
            try:
                result = asyncio.run(
                    query_agent.ainvoke(
                        {"messages": [("user", demand)]},
                        config={"recursion_limit": 100}
                    )
                )
                return {
                    "value": result["messages"][-1].content,
                    "index": demand,
                    "status": "queried"
                }
            except (httpx.RemoteProtocolError, httpx.ConnectTimeout) as e:
                logger.warning(f"连接异常: {demand} - {str(e)}")
                return None
            except Exception as e:
                logger.error(f"未知错误: {demand} - {str(e)}")
                return None

        with ThreadPoolExecutor(max_workers=5) as executor:
            futures = [
                executor.submit(safe_query, demand)
                for demand in state["demandList"]
            ]
            
            results = []
            for future in as_completed(futures):
                res = future.result()
                if res:
                    results.append(res)
                    success_count += 1
                else:
                    error_count += 1

        print(f"✅ 成功获取 {success_count} 条，失败 {error_count} 条")
        
        return {"currentData": results}
    # 数据压缩节点
    async def refine_data_step(state: AgentState):
        from asyncio import create_task, gather

        # 异步处理非缓存数据
        async def process_item(raw_data):
            try:
                refineValue = await refine.ainvoke({"currentData": raw_data["value"]})
                if not refineValue.isNone:
                    value = {
                        "compressed": refineValue.compressed,
                        "mapping": refineValue.mapping,
                        "key": raw_data["index"]
                    }
                    print(f"✅ 压缩完成 {value}")
                    return value
            except Exception as e:
                logger.error(f"压缩失败: {raw_data['index']} - {str(e)}")
            return None

        # 批量执行异步任务
        tasks = [create_task(process_item(item)) for item in state["currentData"]]
        processed = await gather(*tasks)
        
        # 合并结果
        print(f"✅ 压缩完成 {len(processed)} 条数据")
        return {"data": processed}
    # 数据完备性校验节点
    def data_check_step(state: AgentState):
        from collections import Counter

        # 获取需求列表和实际数据键列表
        expected = state["demandList"]
        actual = [item["key"] for item in state["data"]]
        
        # 使用Counter进行精确计数比较
        expected_counter = Counter(expected)
        actual_counter = Counter(actual)
        
        # 计算缺失项（考虑重复需求的情况）
        missing_counter = expected_counter - actual_counter
        missing_elements = list(missing_counter.elements())
        total_retries = state["total_retries"]
        
        # 输出缺失项
        if missing_elements:
            print(f"⚠️ 缺失数据项：{missing_elements}")
            total_retries += 1
            print(f"🔁 尝试重新获取 {total_retries} 次")
            # 判断是否达到最大重试次数
            if total_retries >= 3:
                str = ""
                for item in missing_elements:
                    str += "【" + item + "】"
                return {"demandList": [], "error_tip": f"无法获取{str}，请确认查询条件是否正确"}
            # 未达到最大重试次数，返回缺失项列表
            return {"demandList": missing_elements, "total_retries": total_retries}
        
        
        # 完全匹配时返回空列表
        print("✅ 数据完备性校验通过")
        return {"demandList": []}
    # 数据加工输出节点
    def process_data_step(state: AgentState):
        # 优先处理系统级错误，无法获取数据时返回错误提示
        error_tip = state.get("error_tip")
        if error_tip:
            return {"output": error_tip}

        params = f"""
            现有数据：{state["data"]}\n
            参考模板：{state["template"]}
        """
        print("---------数据加工执行参数：-", params)
        config = {
            "recursion_limit": 100,
        }
        processValue = process.invoke({"messages": [("user", params)]}, config)
        print("---------数据加工执行结果：-", processValue["messages"][-1].content)
        return {"output": processValue["messages"][-1].content}
    
    def end_step(state: AgentState):
        # 增加校验逻辑确保结果有效性
        final_result = state.get("output", "未获取到有效结果")
        # print(f"最终结果：{final_result}")  # ✅ 添加输出确认
        return {"output": final_result, "__end__": True}
    
    from langgraph.graph import StateGraph, START, END
    workflow = StateGraph(AgentState)
    workflow.add_node("generate", generate_step)
    workflow.add_node("inspect", inspect_step)
    workflow.add_node("verify", verify_step)
    workflow.add_node("parallelQuery", parallel_query_step)
    workflow.add_node("refineData", lambda state: asyncio.run(refine_data_step(state)))
    workflow.add_node("dataCheck", data_check_step)
    workflow.add_node("processData", process_data_step)
    workflow.add_node("end", end_step)
    # 边连接逻辑
    workflow.add_edge(START, "generate")
    workflow.add_edge("generate", "inspect")
    workflow.add_edge("inspect", "verify")
    workflow.add_conditional_edges(
        "verify",
        lambda state: "processData" if len(state["demandList"]) == 0 else "parallelQuery"
    )
    workflow.add_edge("parallelQuery", "refineData")
    workflow.add_edge("refineData", "dataCheck")
    workflow.add_conditional_edges(
        "dataCheck",
        lambda state: "processData" if len(state["demandList"]) == 0 else "parallelQuery"
    )
    workflow.add_edge("processData", "end")
    workflow.add_edge("end", END)
    app = workflow.compile()
    config = {
        "recursion_limit": 100,
        "configurable": {"thread_id": "1"},
    }
    inputs = {"input": input}
    response = app.invoke(inputs, config)
    print("返回结果：", response["output"])
    return response["output"]

# main("""三年级学生学习成绩的分析报告
# 第一段：三年级总体情况
# 三年级共有两个班级，每班各有10名学生，总计20名学生。男女比例均衡，各有10名男生和10名女生。学生的平均年龄为8.5岁。
# 第二段：学生总成绩情况
# 学生的总成绩分布在239分至270分之间，平均总成绩为259.1分。各科平均分如下：数学为88.25分，英语为87分，语文为83.85分。共有10名学生的总成绩低于平均分，占比50%。
# 第三段：提高成绩的建议
# 根据线性回归分析，数学、语文、英语三门课程的投入时间与指标值之间存在显著的正相关关系。对于总成绩低于班级平均分的学生，提出以下两条建议以提升成绩：
# 1.	增加学习时间：回归分析显示，每增加一小时的学习时间，数学成绩提高2分，语文成绩提高2.5分，英语成绩提高4分。因此，建议这些学生在数学、语文和英语上适当增加学习时间。
# 优化时间分配：鉴于英语成绩对学习时间的敏感度最高，建议学生优先保证英语学习的时间，同时合理分配数学和语文的学习时间，以实现总分最大化。
# """)




from fastapi import FastAPI, HTTPException
from typing import Union
import logging
# 设置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI(title="agent", version="0.0.1")
# 定义请求模型
class QueryRequest(BaseModel):
    input: str
# 定义响应模型
class QueryResponse(BaseModel):
    status: int
    data: Union[str, list, dict]
    message: str

class ErrorResponse(BaseModel):
    error: str
    message: str
    status: int
    data: None



@app.post("/langchain/getPrompt", response_model=QueryResponse)
def query_api(query: QueryRequest):
    logger.info(f"Received query: {query.input}")
    try:
        res = match_prompt(query.input)
        print("返回结果res：", res)
        logger.info(f"Response: {res}")
        return QueryResponse(status=200, data=res, message="success")
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        raise HTTPException(status_code=500, detail=str(e))
@app.post("/langchain/chat", response_model=QueryResponse)
def query_api(query: QueryRequest):
    logger.info(f"Received query: {query.input}")
    try:
        input = query.input.split("?")[0]
        template = query.input
        if (input == "学生考试成绩分析报告" or input == "班级考试成绩分析报告"):
            template_results = match_template(input)
            if len(template_results) > 0:
                template = template_results[0]["content"]
            if input ==  "学生考试成绩分析报告":
                template = query.input.split("?")[1] + "\n" + template
        res = main(template)
        print("返回结果res：", res)
        logger.info(f"Response: {res}")
        return QueryResponse(status=200, data=res, message="success")
    except Exception as e:
        logger.error(f"Error processing query: {e}")
        raise HTTPException(status_code=500, detail=str(e))




if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        timeout_keep_alive=300,  # 保持连接超时
        limit_max_requests=1000  # 最大并发请求数
    )